package db

import (
	"encoding/json"

	"github.com/aquasecurity/trivy-db/pkg/types"

	bolt "go.etcd.io/bbolt"
	"golang.org/x/xerrors"
)

const (
	vulnerabilityBucket = "vulnerability"
)

func (dbc Config) PutVulnerability(tx *bolt.Tx, cveID string, vuln types.Vulnerability) error {
	bucket, err := tx.CreateBucketIfNotExists([]byte(vulnerabilityBucket))
	if err != nil {
		return xerrors.Errorf("failed to create a bucket: %w", err)
	}
	v, err := json.Marshal(vuln)
	if err != nil {
		return xerrors.Errorf("failed to marshal JSON: %w", err)
	}
	return bucket.Put([]byte(cveID), v)
}

func (dbc Config) GetVulnerability(cveID string) (vuln types.Vulnerability, err error) {
	err = db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket([]byte(vulnerabilityBucket))
		value := bucket.Get([]byte(cveID))
		if err = json.Unmarshal(value, &vuln); err != nil {
			return xerrors.Errorf("failed to marshal JSON: %w", err)
		}
		return nil
	})
	if err != nil {
		return types.Vulnerability{}, xerrors.Errorf("failed to get the vulnerability: %w", err)
	}
	return vuln, nil
}
